`timescale 1ps/1ps
`include "Define.vh"
module MMU(
    input rst,
    input clk,
    input [63:0] satp,
    input is_ecall,
    input is_mret,
    input tlb_flush,
    input maybe_jump,

    input wire if_stall_from_mem,
    output wire if_stall_to_cpu,

    input wire mem_stall_from_mem,
    output wire mem_stall_to_cpu,

    input wire [63:0] pc_from_cpu,
    output wire [63:0] pc_to_mem,

    input wire if_request_from_cpu,
    output wire if_request_to_mem,

    input wire [31:0] inst_from_mem,
    output wire [31:0] inst_to_cpu,

    input wire [63:0] address_from_cpu,
    output wire [63:0] address_to_mem,

    input wire re_mem_from_cpu,
    output wire re_mem_to_mem,

    output wire [63:0] rdata_mem_to_cpu,
    input wire [63:0] rdata_mem_from_mem,

    input wire we_mem_from_cpu,
    output wire we_mem_to_mem,

    input wire [63:0] wdata_mem_from_cpu,
    output wire [63:0] wdata_mem_to_mem,

    input wire [7:0] wmask_mem_from_cpu,
    output wire [7:0] wmask_mem_to_mem,

    input jump,
    input M_mode
);
    wire M_mode_judge = M_mode && (pc_from_cpu >= 64'h80000000 && pc_from_cpu <= 64'h80200000);
    wire [3:0] satp_mode = satp[63:60];
    wire mem_busy_flush = jump;
    localparam IDLE = 3'b000;
    localparam LEVEL2 = 3'b001;
    localparam LEVEL1 = 3'b010;
    localparam LEVEL0 = 3'b011;
    localparam DONE = 3'b100;
    localparam DUMMY = 3'b101;


    logic [63:0] update_phy_pc;
    logic [63:0] tlb_phy_pc;
    logic itlb_update;
    logic itlb_hit;
    TLB itlb(
        .clk(clk),
        .rst(rst | tlb_flush),
        .virt_addr(pc_from_cpu),
        .in_phys_addr(update_phy_pc),
        .update(itlb_update),
        .tlb_flush(tlb_flush),

        .out_phys_addr(tlb_phy_pc),
        .hit(itlb_hit)
    );

    // logic [63:0] update_phy_addr;
    // logic [63:0] tlb_phy_addr;
    // logic dtlb_update;
    // logic dtlb_hit;
    // TLB #(
    //     parameter ADDR_WIDTH = 64,
    //     parameter PPN_WIDTH = 44,
    //     parameter PAGE_SIZE = 4096,
    //     parameter TLB_SIZE = 32
    // ) dtlb(
    //     .clk(clk),
    //     .rst(rst | tlb_flush),
    //     .virt_addr(address_from_cpu),
    //     .in_phys_addr(update_phy_addr),
    //     .update(dtlb_update),
    //     .out_phys_addr(tlb_phy_addr),
    //     .hit(dtlb_hit)
    // );

    reg lock;
    
    reg if_stall_reg;
    reg [63:0] inst_addr;
    reg [31:0] inst;
    reg if_request;
    reg inst_fetch_finish;


    reg [2:0] pc_state;

    logic is_ls = inst_from_mem[6:0]==7'b0000011|inst_from_mem[6:0]==7'b0100011;
    logic [8:0] pc_VA_VPN2 = pc_from_cpu[38:30];
    logic [8:0] pc_VA_VPN1 = pc_from_cpu[29:21];
    logic [8:0] pc_VA_VPN0 = pc_from_cpu[20:12];
    logic [11:0] pc_VA_OFFSET = pc_from_cpu[11:0];

    reg mem_fetch_finish;

    reg mem_stall_reg;

    reg [63:0] address_reg;
    reg re_mem_reg;
    reg [63:0] rdata_reg;

    reg we_mem_reg;
    reg [63:0] wdata_reg;
    reg [7:0] wmask_reg;
    reg [31:0] inst_reg;
    wire [63:0] pte = rdata_mem_from_mem;

    reg inst_busy;
    reg mem_busy;
    reg dirty_inst;
    reg flag;

   always @(posedge clk or negedge rst) begin
        if(rst) begin
            inst_fetch_finish <= 1;
            mem_fetch_finish <= 1;
            pc_state <= IDLE;
            we_mem_reg <= 0;
            re_mem_reg <= 0;
        end    
        else if(satp_mode!=8 ) begin // SATP == 0 或者不是 SV39模式
            inst_fetch_finish <= 1;
            pc_state <= IDLE;
        end
        else if(mem_busy & ~mem_busy_flush) begin
            if_stall_reg <= 1;
            if_request <= 0;
        end
        else if(satp_mode==8)begin  // SV39模式
            // if(mem_busy_flush) begin
            //     mem_busy <= 0;
            //     data_state <= IDLE;
            // end
            inst_fetch_finish <= 0;
            mem_fetch_finish <= 0;
            case (pc_state)
                IDLE: begin
                    flag <= 0;
                    itlb_update <= 0;
                    if(if_request_from_cpu & ~itlb_hit) begin
                        inst_busy <= 1;

                        if_stall_reg <= 1;
                        mem_stall_reg <= 0;

                        if_request <= 0;
                        address_reg <= {8'b0,satp[43:0],pc_VA_VPN2,3'b0};
                        we_mem_reg <= 0;

                        inst_reg <= inst_from_mem;
                        pc_state <= LEVEL2;
                        re_mem_reg <= 1;
                    end
                    else if((inst_from_mem[6:0]==7'b0000011|inst_from_mem[6:0]==7'b0100011) & ~if_stall_from_mem & ~jump) begin // hit 时取到了 l/s 指令
                        mem_busy <= 1;
                        if_stall_reg <= 1;
                        address_reg <= address_from_cpu;
                    end
                    else begin
                        if_stall_reg <= 0;
                        mem_stall_reg <= 0;
                    end
                end

                DONE: begin
                    if(flag==0) begin
                        itlb_update <= 1;
                        flag <= 1;
                    end
                    else begin
                        itlb_update <= 0;
                    end
                    update_phy_pc <= inst_addr;

                    if(!if_stall_from_mem) begin
                        flag <= 0;
                        if_request <= 0;
                        inst <= inst_from_mem;
                        if_stall_reg <= 0;
                        mem_stall_reg <= 0;
                        pc_state <= IDLE;
                        inst_busy <= 0;
                        if(inst_from_mem[6:0]==7'b0000011|inst_from_mem[6:0]==7'b0100011) begin
                            if(dirty_inst) begin
                                mem_busy <= 0;
                                dirty_inst <= 0;
                            end else begin
                                mem_busy <= 1;
                            end
                            address_reg <= address_from_cpu;
                        end else begin
                            dirty_inst <= 0;
                        end
                    end
                end

                LEVEL2: begin
                    if(!dirty_inst) begin
                        dirty_inst <= jump;
                    end

                    if(mem_stall_from_mem==0) begin

                        re_mem_reg <= 0;
                        if(pte[3:1]==3'b000 && pte[0]==1) begin
                            address_reg <= {8'b0,pte[53:10],pc_VA_VPN1,3'b0};
                            pc_state <= LEVEL1;
                            // re_mem_reg <= 1;
                        end
                        else begin
                            if_request <= 1;
                            re_mem_reg <= 0;
                            inst_addr <=  {8'b0,pte[53:28],pc_VA_VPN1,pc_VA_VPN0,pc_VA_OFFSET};
                            pc_state <= DONE;
                            
                        end
                    end
                end

                LEVEL1: begin
                    if(!dirty_inst) begin
                        dirty_inst <= jump;
                    end
                    
                    if(re_mem_reg==0) begin
                        re_mem_reg <= 1;
                    end
                    else if(re_mem_reg == 1 & mem_stall_from_mem==0) begin

                        re_mem_reg <= 0;
                        if(pte[3:1]==3'b000 && pte[0]==1) begin
                            address_reg <= {8'b0,pte[53:10],pc_VA_VPN0,3'b0};
                            pc_state <= LEVEL0;
                            // re_mem_reg <= 1;
                        end
                        else begin
                            if_request <= 1;
                            re_mem_reg <= 0;
                            inst_addr <=  {8'b0,pte[53:19],pc_VA_VPN0,pc_VA_OFFSET};
                            pc_state <= DONE;
                        end
                    end
                end

                LEVEL0: begin
                    if(!dirty_inst) begin
                        dirty_inst <= jump;
                    end

                    if(re_mem_reg==0) begin
                        re_mem_reg <= 1;
                    end
                    else if(re_mem_reg==1 && mem_stall_from_mem==0) begin
                        if_request <= 1;
                        re_mem_reg <= 0;
                        inst_addr <=  {8'b0,pte[53:10],pc_VA_OFFSET};
                        pc_state <= DONE;
                    end
                end
                DUMMY:begin
                    pc_state <= IDLE;
                    if_stall_reg <= 1;
                end
                default:begin
                end
            endcase
        end
    end

    reg [2:0] data_state;
    logic [8:0] data_VA_VPN2 = address_from_cpu[38:30];
    logic [8:0] data_VA_VPN1 = address_from_cpu[29:21];  
    logic [8:0] data_VA_VPN0 = address_from_cpu[20:12];
    logic [11:0] data_VA_OFFSET = address_from_cpu[11:0];

    reg do_write;
    reg do_read;
    reg flag2 = 0;
    // data
    always @(posedge clk or negedge rst) begin

        if(rst) begin
            inst_fetch_finish <= 1;
            mem_fetch_finish <= 1;
            data_state <= IDLE;
        end
        else if(satp_mode!=8 ) begin
            // inst_fetch_finish <= 1;
            mem_fetch_finish <= 1;
            data_state <= IDLE;
        end
        else if(mem_busy) begin  // SV39模式
            inst_fetch_finish <= 0;
            
            case (data_state)
                IDLE: begin
                    if_stall_reg <= 1;
                    if_request <= 0;
                    we_mem_reg <= 0;
                    re_mem_reg <= 0;
                    if(re_mem_from_cpu|we_mem_from_cpu) begin // 等到 s/l 指令到MEM阶段
                        mem_stall_reg <= 1;
                        do_write <= we_mem_from_cpu;
                        do_read <= re_mem_from_cpu;
                        wdata_reg <= wdata_mem_from_cpu;
                        wmask_reg <= wmask_mem_from_cpu;

                        address_reg <= {8'b0,satp[43:0],data_VA_VPN2,3'b0};
                        data_state <= LEVEL2;
                        re_mem_reg <= 1;
                    end
                    else if(dirty_inst | jump) begin // jump 之后多取了 s/l 指令
                        pc_state <= DUMMY;
                        dirty_inst <= 0;
                        mem_stall_reg <= 0;
                        mem_busy <= 0;
                    end
                    else begin
                        mem_fetch_finish <= 0;
                        if(!dirty_inst) begin
                            dirty_inst <= jump;
                        end
                    end
                end
                DONE: begin
                    // we_mem_reg <= 0;
                    mem_stall_reg <= mem_stall_from_mem;
                    if(!mem_stall_from_mem) begin
                        we_mem_reg <= 0;
                        address_reg <= pc_from_cpu;
                        data_state<= IDLE;
                        mem_busy <= 0;
                        rdata_reg <= rdata_mem_from_mem;
                        re_mem_reg <= 0;
                    end
                    else begin
                        // re_mem_reg <= 0;
                    end
                end
                LEVEL2: begin
                    if(!dirty_inst) begin
                        dirty_inst <= jump;
                    end
                    if(mem_stall_from_mem==0) begin
                        // re_mem_reg <= 0;
                        if(pte[3:1]==3'b000 && pte[0]==1) begin
                            address_reg <= {8'b0,pte[53:10],data_VA_VPN1,3'b0};
                            data_state <= LEVEL1;
                            re_mem_reg <= 1;
                        end
                        else begin  //结束
                            re_mem_reg <= do_read;
                            we_mem_reg <= do_write;
                            address_reg <= {8'b0,pte[53:28],data_VA_VPN1,data_VA_VPN0,data_VA_OFFSET};
                            data_state <= DONE;
                        end
                    end
                    else begin
                    end
                end
                LEVEL1: begin
                    if(!dirty_inst) begin
                        dirty_inst <= jump;
                    end
                    if(mem_stall_from_mem==0) begin
                        if(pte[3:1]==3'b000 && pte[0]==1) begin
                            address_reg <= {8'b0,pte[53:10],data_VA_VPN0,3'b0};
                            data_state <= LEVEL0;
                            re_mem_reg <= 1;
                        end
                        else begin
                            re_mem_reg <= do_read;
                            we_mem_reg <= do_write;
                            address_reg <= {8'b0,pte[53:19],data_VA_VPN0,data_VA_OFFSET};
                            data_state <= DONE;
                        end
                    end
                    else begin
                        if(!dirty_inst) begin
                            dirty_inst <= jump;
                        end
                    end
                end
                LEVEL0: begin
                    if(!dirty_inst) begin
                        dirty_inst <= jump;
                    end
                    if(mem_stall_from_mem==0) begin
                        re_mem_reg <= do_read;
                        we_mem_reg <= do_write;
                        address_reg <=  {8'b0,pte[53:10],data_VA_OFFSET};
                        data_state <= DONE;
                    end
                    else begin
                    end
                end
                default: begin

                end
            endcase
        end
    end

assign if_stall_to_cpu = inst_fetch_finish?if_stall_from_mem:(if_stall_reg | if_stall_from_mem | maybe_jump);
assign pc_to_mem = (inst_fetch_finish|M_mode_judge)?pc_from_cpu:(itlb_hit?tlb_phy_pc:inst_addr);
assign if_request_to_mem = ((inst_fetch_finish|itlb_hit)?if_request_from_cpu:if_request) & (pc_to_mem <= 64'hFFFFFFE000000000);
assign inst_to_cpu = (inst_fetch_finish|itlb_hit)?inst_from_mem:inst;

assign mem_stall_to_cpu=mem_fetch_finish?mem_stall_from_mem:(mem_stall_reg | mem_stall_from_mem | ((re_mem_from_cpu | we_mem_from_cpu) & mem_busy));
assign address_to_mem=(mem_fetch_finish|M_mode_judge)?address_from_cpu:address_reg;
assign re_mem_to_mem=mem_fetch_finish?re_mem_from_cpu:re_mem_reg;
assign rdata_mem_to_cpu=mem_fetch_finish?rdata_mem_from_mem:rdata_reg;
assign we_mem_to_mem=mem_fetch_finish?we_mem_from_cpu:we_mem_reg;
assign wdata_mem_to_mem=(mem_fetch_finish|M_mode_judge)?wdata_mem_from_cpu:wdata_reg;
assign wmask_mem_to_mem=(mem_fetch_finish|M_mode_judge)?wmask_mem_from_cpu:wmask_reg;

endmodule